#!/usr/bin/env python3
"""
性能对比和可视化脚本
"""

import os
import json
import argparse
from pathlib import Path
from typing import Dict, List
import matplotlib.pyplot as plt  # type: ignore
import matplotlib  # type: ignore
matplotlib.use('Agg')  # 非交互式后端
import numpy as np
import pandas as pd


def load_metrics(metrics_file: str) -> Dict:
    """加载性能指标文件"""
    with open(metrics_file, 'r') as f:
        return json.load(f)


def compare_summary(pytorch_metrics: Dict, triton_metrics: Dict):
    """对比摘要统计"""
    pytorch_summary = pytorch_metrics['summary']
    triton_summary = triton_metrics['summary']
    
    print("\n" + "="*80)
    print("Performance Comparison Summary")
    print("="*80)
    
    metrics_to_compare = [
        ('avg_loss', 'Average Loss', '.4f'),
        ('avg_throughput_samples', 'Throughput (samples/sec)', '.2f'),
        ('avg_throughput_tokens', 'Throughput (tokens/sec)', '.2f'),
        ('avg_memory_allocated', 'Memory Allocated (GB)', '.2f'),
        ('avg_step_time', 'Step Time (sec)', '.4f'),
        ('total_time', 'Total Time (sec)', '.2f'),
    ]
    
    print(f"\n{'Metric':<30} {'PyTorch':<15} {'Triton':<15} {'Speedup':<15}")
    print("-"*80)
    
    for key, name, fmt in metrics_to_compare:
        pytorch_val = pytorch_summary.get(key, 0)
        triton_val = triton_summary.get(key, 0)
        
        if pytorch_val > 0:
            if 'time' in key.lower():
                # 对于时间指标，speedup = pytorch/triton (越小越好)
                speedup = pytorch_val / triton_val if triton_val > 0 else 0
                speedup_str = f"{speedup:.2f}x"
            elif 'throughput' in key.lower():
                # 对于吞吐量，speedup = triton/pytorch (越大越好)
                speedup = triton_val / pytorch_val if pytorch_val > 0 else 0
                speedup_str = f"{speedup:.2f}x"
            else:
                speedup_str = "-"
        else:
            speedup_str = "-"
        
        print(f"{name:<30} {pytorch_val:<15{fmt}} {triton_val:<15{fmt}} {speedup_str:<15}")
    
    print("="*80)


def plot_throughput_comparison(pytorch_metrics: Dict, triton_metrics: Dict, output_dir: str):
    """绘制吞吐量对比图"""
    pytorch_throughput = pytorch_metrics['detailed_metrics']['throughput_tokens']
    triton_throughput = triton_metrics['detailed_metrics']['throughput_tokens']
    
    steps = range(len(pytorch_throughput))
    
    plt.figure(figsize=(12, 6))
    plt.plot(steps, pytorch_throughput, label='PyTorch', alpha=0.7)
    plt.plot(steps, triton_throughput, label='Triton', alpha=0.7)
    plt.xlabel('Training Step')
    plt.ylabel('Throughput (tokens/sec)')
    plt.title('Training Throughput Comparison')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    
    output_path = Path(output_dir) / "throughput_comparison.png"
    plt.savefig(output_path, dpi=300)
    print(f"✓ Saved throughput comparison to {output_path}")
    plt.close()


def plot_memory_comparison(pytorch_metrics: Dict, triton_metrics: Dict, output_dir: str):
    """绘制内存使用对比图"""
    pytorch_memory = pytorch_metrics['detailed_metrics']['memory_allocated']
    triton_memory = triton_metrics['detailed_metrics']['memory_allocated']
    
    steps = range(len(pytorch_memory))
    
    plt.figure(figsize=(12, 6))
    plt.plot(steps, pytorch_memory, label='PyTorch', alpha=0.7)
    plt.plot(steps, triton_memory, label='Triton', alpha=0.7)
    plt.xlabel('Training Step')
    plt.ylabel('Memory Allocated (GB)')
    plt.title('Memory Usage Comparison')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    
    output_path = Path(output_dir) / "memory_comparison.png"
    plt.savefig(output_path, dpi=300)
    print(f"✓ Saved memory comparison to {output_path}")
    plt.close()


def plot_loss_comparison(pytorch_metrics: Dict, triton_metrics: Dict, output_dir: str):
    """绘制loss对比图"""
    pytorch_loss = pytorch_metrics['detailed_metrics']['loss']
    triton_loss = triton_metrics['detailed_metrics']['loss']
    
    steps = range(len(pytorch_loss))
    
    plt.figure(figsize=(12, 6))
    plt.plot(steps, pytorch_loss, label='PyTorch', alpha=0.7)
    plt.plot(steps, triton_loss, label='Triton', alpha=0.7)
    plt.xlabel('Training Step')
    plt.ylabel('Loss')
    plt.title('Training Loss Comparison')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    
    output_path = Path(output_dir) / "loss_comparison.png"
    plt.savefig(output_path, dpi=300)
    print(f"✓ Saved loss comparison to {output_path}")
    plt.close()


def generate_comparison_report(pytorch_metrics: Dict, triton_metrics: Dict, output_dir: str):
    """生成对比报告"""
    output_path = Path(output_dir) / "comparison_report.md"
    
    pytorch_summary = pytorch_metrics['summary']
    triton_summary = triton_metrics['summary']
    
    with open(output_path, 'w') as f:
        f.write("# Triton vs PyTorch Performance Comparison\n\n")
        # Save console summary as a fenced code block
        f.write("```\n")
        f.write("="*80 + "\n")
        f.write("Performance Comparison Summary\n")
        f.write("="*80 + "\n\n")
        console_metrics = [
            ("Average Loss", pytorch_summary.get("avg_loss", 0), triton_summary.get("avg_loss", 0), "loss"),
            ("Throughput (samples/sec)", pytorch_summary.get("avg_throughput_samples", 0), triton_summary.get("avg_throughput_samples", 0), "throughput"),
            ("Throughput (tokens/sec)", pytorch_summary.get("avg_throughput_tokens", 0), triton_summary.get("avg_throughput_tokens", 0), "throughput"),
            ("Memory Allocated (GB)", pytorch_summary.get("avg_memory_allocated", 0), triton_summary.get("avg_memory_allocated", 0), "memory"),
            ("Step Time (sec)", pytorch_summary.get("avg_step_time", 0), triton_summary.get("avg_step_time", 0), "time"),
            ("Total Time (sec)", pytorch_summary.get("total_time", 0), triton_summary.get("total_time", 0), "time"),
        ]
        for name, pt_val, tr_val, mtype in console_metrics:
            if mtype == "throughput":
                sp = f"{(tr_val / pt_val if pt_val > 0 else 0):.2f}x"
            elif mtype == "time":
                sp = f"{(pt_val / tr_val if tr_val > 0 else 0):.2f}x"
            else:
                sp = "-"
            # select precision
            prec = 4 if name in ["Average Loss", "Step Time (sec)"] else 2
            f.write(f"{name:<30} {pt_val:<15.{prec}f} {tr_val:<15.{prec}f} {sp:<15}\n")
        f.write("="*80 + "\n")
        f.write("```\n\n")
        f.write("## Summary\n\n")
        f.write("| Metric | PyTorch | Triton | Improvement |\n")
        f.write("|--------|---------|--------|-------------|\n")
        
        # 吞吐量
        pt_throughput = pytorch_summary['avg_throughput_tokens']
        tr_throughput = triton_summary['avg_throughput_tokens']
        improvement = (tr_throughput / pt_throughput - 1) * 100 if pt_throughput > 0 else 0
        f.write(f"| Throughput (tokens/sec) | {pt_throughput:.2f} | {tr_throughput:.2f} | {improvement:+.1f}% |\n")
        
        # 内存
        pt_memory = pytorch_summary['avg_memory_allocated']
        tr_memory = triton_summary['avg_memory_allocated']
        improvement = (1 - tr_memory / pt_memory) * 100 if pt_memory > 0 else 0
        f.write(f"| Memory (GB) | {pt_memory:.2f} | {tr_memory:.2f} | {improvement:+.1f}% |\n")
        
        # 时间
        pt_time = pytorch_summary['avg_step_time']
        tr_time = triton_summary['avg_step_time']
        improvement = (1 - tr_time / pt_time) * 100 if pt_time > 0 else 0
        f.write(f"| Step Time (sec) | {pt_time:.4f} | {tr_time:.4f} | {improvement:+.1f}% |\n")
        
        f.write("\n## Plots\n\n")
        f.write("![Throughput Comparison](throughput_comparison.png)\n\n")
        f.write("![Memory Comparison](memory_comparison.png)\n\n")
        f.write("![Loss Comparison](loss_comparison.png)\n\n")
    
    print(f"✓ Saved comparison report to {output_path}")


def load_config(config_file: str) -> Dict:
    """加载配置文件"""
    with open(config_file, 'r') as f:
        return json.load(f)



# 用于生成比较文件的两个文件分别来自pytorch.json和triton.json配置文件的output_dir
# 生成的比较文件存放在pytorch.json配置文件的output_dir的父目录的comparison目录下的子目录中
def main():
    parser = argparse.ArgumentParser(description="Compare PyTorch and Triton performance")
    
    # 使用固定的配置文件路径
    pytorch_config_path = "../../configs/pytorch.json"
    triton_config_path = "../../configs/triton.json"
    
    # 从配置文件加载output_dir并构建metrics路径
    pytorch_config = load_config(pytorch_config_path)
    triton_config = load_config(triton_config_path)
    
    # 使用配置文件中的output_dir构建metrics文件路径
    pytorch_output_dir = pytorch_config['training']['output_dir']
    pytorch_metrics = os.path.join(pytorch_output_dir, "metrics.json")
    
    triton_output_dir = triton_config['training']['output_dir']
    triton_metrics = os.path.join(triton_output_dir, "metrics.json")
    
    # 使用PyTorch配置的output_dir的父目录作为比较结果保存路径
    comparison_output_dir = os.path.dirname(pytorch_output_dir)
    comparison_output_dir = os.path.join(comparison_output_dir, "comparison", "triton_all")
    
    # 创建输出目录
    output_dir = Path(comparison_output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 加载指标
    print("Loading metrics...")
    pytorch_metrics_data = load_metrics(pytorch_metrics)
    triton_metrics_data = load_metrics(triton_metrics)
    
    # 打印对比摘要
    compare_summary(pytorch_metrics_data, triton_metrics_data)
    
    # 生成可视化
    print("\nGenerating plots...")
    plot_throughput_comparison(pytorch_metrics_data, triton_metrics_data, output_dir)
    plot_memory_comparison(pytorch_metrics_data, triton_metrics_data, output_dir)
    plot_loss_comparison(pytorch_metrics_data, triton_metrics_data, output_dir)
    
    # 生成报告
    print("\nGenerating report...")
    generate_comparison_report(pytorch_metrics_data, triton_metrics_data, output_dir)
    
    print("\n" + "="*80)
    print("Comparison completed!")
    print(f"Results saved to {output_dir}")
    print("="*80)


if __name__ == "__main__":
    main()

